{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "ExecuteTime": {
     "end_time": "2025-07-10T06:31:22.657795Z",
     "start_time": "2025-07-10T06:31:22.650288Z"
    }
   },
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "batch_size = 2\n",
    "seq_len = 5\n",
    "d_model = 512  # 词向量的长度\n",
    "vocab_size = 100\n",
    "\n",
    "# 生成两个句子，每个句子有5个词\n",
    "# 随机生成0-100的数字，形状为(batch_size, seq_len)，表示batch_size条sequence，每条sequence的长度为seq_len\n",
    "inputs = torch.randint(0, vocab_size, (batch_size, seq_len))\n",
    "inputs"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[61, 81, 31, 40, 17],\n",
       "        [ 4, 95, 48, 24, 12]])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 23
  },
  {
   "cell_type": "code",
   "id": "d75a5f086f4eea9c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T06:31:22.704768Z",
     "start_time": "2025-07-10T06:31:22.701215Z"
    }
   },
   "source": [
    "# 词嵌入层\n",
    "# 两个句子，每个句子有5个词，每个词的词向量长度为512\n",
    "embedding = nn.Embedding(vocab_size, d_model)\n",
    "embeddings = embedding(inputs)\n",
    "embeddings.shape"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 5, 512])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 24
  },
  {
   "cell_type": "code",
   "id": "b78fcbd2c1eeab9e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T06:31:22.750781Z",
     "start_time": "2025-07-10T06:31:22.746023Z"
    }
   },
   "source": [
    "import math\n",
    "\n",
    "# 位置编码\n",
    "# 位置编码只跟位置有关，和具体位置上输入的数据无关\n",
    "pe = torch.zeros(seq_len, d_model)\n",
    "position = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1)\n",
    "div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))\n",
    "pe[:, 0::2] = torch.sin(position * div_term)\n",
    "pe[:, 1::2] = torch.cos(position * div_term)\n",
    "pe"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.0000e+00,  1.0000e+00,  0.0000e+00,  ...,  1.0000e+00,\n",
       "          0.0000e+00,  1.0000e+00],\n",
       "        [ 8.4147e-01,  5.4030e-01,  8.2186e-01,  ...,  1.0000e+00,\n",
       "          1.0366e-04,  1.0000e+00],\n",
       "        [ 9.0930e-01, -4.1615e-01,  9.3641e-01,  ...,  1.0000e+00,\n",
       "          2.0733e-04,  1.0000e+00],\n",
       "        [ 1.4112e-01, -9.8999e-01,  2.4509e-01,  ...,  1.0000e+00,\n",
       "          3.1099e-04,  1.0000e+00],\n",
       "        [-7.5680e-01, -6.5364e-01, -6.5717e-01,  ...,  1.0000e+00,\n",
       "          4.1465e-04,  1.0000e+00]])"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 25
  },
  {
   "cell_type": "code",
   "id": "220f121c2cc2e2ef",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T06:31:22.800216Z",
     "start_time": "2025-07-10T06:31:22.797321Z"
    }
   },
   "source": [
    "# 位置编码+词向量\n",
    "decoder_inputs = embeddings + pe\n",
    "decoder_inputs.shape"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 5, 512])"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 26
  },
  {
   "cell_type": "code",
   "id": "b87ee8e199a33508",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T06:31:22.873746Z",
     "start_time": "2025-07-10T06:31:22.868892Z"
    }
   },
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "class MultiHeadAttention(nn.Module):\n",
    "\n",
    "    def __init__(self, embed_dim: int, attn_dim: int, output_dim: int, num_heads: int):\n",
    "        super().__init__()\n",
    "\n",
    "        self.embed_dim = embed_dim\n",
    "        self.attn_dim = attn_dim\n",
    "        self.output_dim = output_dim\n",
    "        self.num_heads = num_heads\n",
    "        self.head_dim = attn_dim // num_heads  # //表示向下取整，attn_dim是head_dim的整数倍\n",
    "\n",
    "        # QKV投影层：从输入维度映射到内部维度\n",
    "        # projection\n",
    "        self.q_proj = nn.Linear(embed_dim, self.attn_dim)\n",
    "        self.k_proj = nn.Linear(embed_dim, self.attn_dim)\n",
    "        self.v_proj = nn.Linear(embed_dim, self.attn_dim)\n",
    "\n",
    "        # 输出投影层：从内部维度映射到输出维度\n",
    "        self.out_proj = nn.Linear(self.attn_dim, self.output_dim)\n",
    "\n",
    "    def forward(self, q_x, k_x, v_x, mask=None):\n",
    "        \"\"\"\n",
    "        输入: [batch_size, seq_len, embed_dim]\n",
    "        返回: [batch_size, seq_len, output_dim]\n",
    "        \"\"\"\n",
    "        batch_size, q_seq_len, embed_dim = q_x.shape\n",
    "        batch_size, k_seq_len, embed_dim = k_x.shape\n",
    "\n",
    "        # 投影到QKV空间\n",
    "        q = self.q_proj(q_x)  # [batch_size, seq_len, attn_dim]\n",
    "        k = self.k_proj(k_x)  # [batch_size, seq_len, attn_dim]\n",
    "        v = self.v_proj(v_x)  # [batch_size, seq_len, attn_dim]\n",
    "\n",
    "        # [batch_size, seq_len, num_heads, head_dim]\n",
    "        # 分割多头 [batch_size, num_heads, seq_len, head_dim]\n",
    "        q = q.view(batch_size, q_seq_len, self.num_heads, self.head_dim).transpose(1, 2)\n",
    "        k = k.view(batch_size, k_seq_len, self.num_heads, self.head_dim).transpose(1, 2)\n",
    "        v = v.view(batch_size, k_seq_len, self.num_heads, self.head_dim).transpose(1, 2)\n",
    "\n",
    "        # 计算注意力得分\n",
    "        # q   [batch_size, num_heads, seq_len, head_dim]\n",
    "        # k.T [batch_size, num_heads, head_dim, seq_len]\n",
    "        # q @ k.T 形状: [batch_size, num_heads, seq_len, seq_len]\n",
    "        attn_scores = torch.matmul(q, k.transpose(-2, -1))\n",
    "\n",
    "        # 缩放因子：防止乘积过大\n",
    "        d_k = k.size(-1)\n",
    "        attn_scores = attn_scores / torch.sqrt(torch.tensor(d_k))\n",
    "\n",
    "        if mask is not None:\n",
    "            attn_scores = attn_scores.masked_fill(mask == 0, float('-inf'))\n",
    "\n",
    "        # 计算注意力权重\n",
    "        attn_weights = torch.softmax(attn_scores, dim=-1)\n",
    "\n",
    "        # 计算注意力输出\n",
    "        # attn_weights [batch_size, num_heads, seq_len, seq_len]\n",
    "        # v            [batch_size, num_heads, seq_len, head_dim]\n",
    "        # [batch_size, num_heads, seq_len, head_dim]\n",
    "        attn_out = torch.matmul(attn_weights, v)\n",
    "\n",
    "        # 合并多头 [batch_size, seq_len, attn_dim]\n",
    "\n",
    "        # [batch_size, seq_len, num_heads, head_dim]\n",
    "        attn_out = attn_out.transpose(1, 2).reshape(batch_size, q_seq_len, self.attn_dim)\n",
    "\n",
    "        # 投影到输出空间\n",
    "        return self.out_proj(attn_out)"
   ],
   "outputs": [],
   "execution_count": 27
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T06:31:23.054932Z",
     "start_time": "2025-07-10T06:31:23.048305Z"
    }
   },
   "cell_type": "code",
   "source": [
    "attn = MultiHeadAttention(embed_dim=d_model, attn_dim=d_model, output_dim=d_model, num_heads=2)\n",
    "\n",
    "mask = torch.tril(torch.ones(seq_len, seq_len))\n",
    "\n",
    "decoder_mask_attn_outputs = attn(q_x=decoder_inputs, k_x=decoder_inputs, v_x=decoder_inputs, mask=mask)\n",
    "\n",
    "print(decoder_inputs.shape)\n",
    "print(decoder_mask_attn_outputs.shape)"
   ],
   "id": "eee8fb2ee0dd3cb2",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 5, 512])\n",
      "torch.Size([2, 5, 512])\n"
     ]
    }
   ],
   "execution_count": 28
  },
  {
   "cell_type": "code",
   "id": "7cadc119cd37aaa3",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T06:31:23.169005Z",
     "start_time": "2025-07-10T06:31:23.166509Z"
    }
   },
   "source": [
    "## 残差连接\n",
    "decoder_mask_attn_add_outputs = decoder_inputs + decoder_mask_attn_outputs\n",
    "decoder_mask_attn_add_outputs.shape"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 5, 512])"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 29
  },
  {
   "cell_type": "code",
   "id": "fa842128f70b648e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T06:31:23.200796Z",
     "start_time": "2025-07-10T06:31:23.197662Z"
    }
   },
   "source": [
    "## 层归一化\n",
    "layer_norm_1 = nn.LayerNorm(d_model)\n",
    "decoder_mask_attn_add_norm_outputs = layer_norm_1(decoder_mask_attn_add_outputs)\n",
    "decoder_mask_attn_add_norm_outputs.shape"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 5, 512])"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 30
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T06:31:23.228124Z",
     "start_time": "2025-07-10T06:31:23.221966Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 模拟Encoder的输出（2, 4, 512）\n",
    "encoder_outputs = torch.randn(batch_size, 4, d_model)\n",
    "\n",
    "cross_attn = MultiHeadAttention(embed_dim=d_model, attn_dim=d_model, output_dim=d_model, num_heads=2)\n",
    "decoder_mask_attn_add_norm_cross_outputs = cross_attn(q_x=decoder_mask_attn_add_norm_outputs, k_x=encoder_outputs,\n",
    "                                                      v_x=encoder_outputs, mask=None)"
   ],
   "id": "a03bdf06262a5c0",
   "outputs": [],
   "execution_count": 31
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T06:31:23.246456Z",
     "start_time": "2025-07-10T06:31:23.243515Z"
    }
   },
   "cell_type": "code",
   "source": [
    "## 残差连接\n",
    "decoder_mask_attn_add_norm_cross_add_outputs = decoder_mask_attn_add_norm_outputs + decoder_mask_attn_add_norm_cross_outputs\n",
    "\n",
    "## 层归一化\n",
    "layer_norm_2 = nn.LayerNorm(d_model)\n",
    "decoder_mask_attn_add_norm_cross_add_norm_outputs = layer_norm_2(decoder_mask_attn_add_norm_cross_add_outputs)\n",
    "decoder_mask_attn_add_norm_cross_add_norm_outputs.shape"
   ],
   "id": "56864fb07f7c30bc",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 5, 512])"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 32
  },
  {
   "cell_type": "code",
   "id": "a004b60ec103d66a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T06:31:23.270201Z",
     "start_time": "2025-07-10T06:31:23.259415Z"
    }
   },
   "source": [
    "## Feed Forward，两个线性层，最后输出维度不变\n",
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, d_model, d_ff):\n",
    "        super().__init__()\n",
    "        self.fc1 = nn.Linear(d_model, d_ff)\n",
    "        self.fc2 = nn.Linear(d_ff, d_model)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.fc2(torch.relu(self.fc1(x)))\n",
    "\n",
    "\n",
    "feed_forward = FeedForward(d_model=d_model, d_ff=2048)\n",
    "\n",
    "decoder_mask_attn_add_norm_cross_add_norm_feed_outputs = feed_forward(decoder_mask_attn_add_norm_cross_add_norm_outputs)\n",
    "decoder_mask_attn_add_norm_cross_add_norm_feed_outputs.shape"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 5, 512])"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 33
  },
  {
   "cell_type": "code",
   "id": "1d01b7ff51008594",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T06:31:23.292505Z",
     "start_time": "2025-07-10T06:31:23.288385Z"
    }
   },
   "source": [
    "# 再来一次残差连接和层归一化，得到最终的Encoder输出\n",
    "layer_norm_3 = nn.LayerNorm(d_model)\n",
    "decoder_outputs = layer_norm_3(\n",
    "    decoder_mask_attn_add_norm_cross_add_norm_feed_outputs + decoder_mask_attn_add_norm_cross_add_norm_outputs)\n",
    "decoder_outputs.shape"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 5, 512])"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 34
  },
  {
   "cell_type": "code",
   "id": "16fba6d2199cee2b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T06:31:23.418118Z",
     "start_time": "2025-07-10T06:31:23.304223Z"
    }
   },
   "source": [
    "# 定义Decoder\n",
    "class Decoder(nn.Module):\n",
    "    def __init__(self, d_model, d_ff, num_heads, num_decoder_layers):\n",
    "        super().__init__()\n",
    "        self.layers = nn.ModuleList([DecoderLayer(d_model, d_ff, num_heads) for _ in range(num_decoder_layers)])\n",
    "\n",
    "    def forward(self, decoder_inputs, encoder_outputs, mask=None):\n",
    "        for layer in self.layers:\n",
    "            decoder_inputs = layer(decoder_inputs, encoder_outputs, mask=mask)\n",
    "        return decoder_inputs\n",
    "\n",
    "\n",
    "# 定义DecoderLayer\n",
    "class DecoderLayer(nn.Module):\n",
    "    def __init__(self, d_model, d_ff, num_heads):\n",
    "        super().__init__()\n",
    "        self.masked_mha = MultiHeadAttention(embed_dim=d_model, attn_dim=d_model, output_dim=d_model,\n",
    "                                             num_heads=num_heads)\n",
    "        self.cross_mha = MultiHeadAttention(embed_dim=d_model, attn_dim=d_model, output_dim=d_model,\n",
    "                                            num_heads=num_heads)\n",
    "        self.ff = FeedForward(d_model, d_ff)\n",
    "        self.layer_norm1 = nn.LayerNorm(d_model)\n",
    "        self.layer_norm2 = nn.LayerNorm(d_model)\n",
    "        self.layer_norm3 = nn.LayerNorm(d_model)\n",
    "\n",
    "    def forward(self, decoder_inputs, encoder_outputs, mask=None):\n",
    "\n",
    "        # 1.masked多头注意力\n",
    "        # 2.残差连接和层归一化\n",
    "        x = self.layer_norm1(decoder_inputs + self.masked_mha(q_x=decoder_inputs, k_x=decoder_inputs, v_x=decoder_inputs, mask=mask))\n",
    "\n",
    "        # 3.cross多头注意力\n",
    "        # 4.残差连接和层归一化\n",
    "        x = self.layer_norm2(x + self.cross_mha(q_x=x, k_x=encoder_outputs, v_x=encoder_outputs))\n",
    "\n",
    "        # 5.Feed Forward\n",
    "        # 6.残差连接和层归一化\n",
    "        return self.layer_norm3(x + self.ff(x))\n",
    "\n",
    "\n",
    "decoder = Decoder(d_model=512, d_ff=2048, num_heads=8, num_decoder_layers=6)\n",
    "\n",
    "mask = torch.tril(torch.ones(seq_len, seq_len))\n",
    "decoder_outputs = decoder(decoder_inputs=decoder_inputs, encoder_outputs=encoder_outputs, mask=mask)\n",
    "decoder_outputs.shape"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 5, 512])"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 35
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T06:33:37.134249Z",
     "start_time": "2025-07-10T06:33:37.110429Z"
    }
   },
   "cell_type": "code",
   "source": [
    "fc = nn.Linear(d_model, vocab_size)\n",
    "decoder_outputs_vocab = fc(decoder_outputs)\n",
    "decoder_outputs_vocab.shape"
   ],
   "id": "1f9f3d935328ae77",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 5, 100])"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 36
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T06:36:50.860566Z",
     "start_time": "2025-07-10T06:36:50.855429Z"
    }
   },
   "cell_type": "code",
   "source": [
    "outputs = torch.softmax(decoder_outputs_vocab, dim=-1)\n",
    "outputs.shape"
   ],
   "id": "19264645e4ea0b0d",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 5, 100])"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 42
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
